﻿#include "Processor.h"


CProcessor::CProcessor(void)
{
	FileStorage fs("config.xml", FileStorage::READ);
		
	fs["DATASET_PATH"] >> this->DATASET_PATH;

	fs["NUMBER_IMAGES"] >> this->NUMBER_IMAGES;

	fs["XML_PATH"] >> this->XML_PATH;

	fs["VOCABULARY_PATH"] >> this->VOCABULARY_PATH;

	fs["DICTIONARY_SIZE"] >> this->DICTIONARY_SIZE;

	this->FILE_LIST = GetFileList(fs["FILE_LIST_PATH"]);

	fs.release();
}

CProcessor::~CProcessor(void)
{

}

void CProcessor::Trainer()
{
	//BEGIN: Load dữ liệu ảnh
	vector<Mat> imgs;

	for (int i = 0; i < NUMBER_IMAGES; i++)
	{
		imgs.push_back(imread(DATASET_PATH + FILE_LIST[i] + ".jpg", CV_LOAD_IMAGE_GRAYSCALE));

	}
	//END: Load dữ liệu ảnh
	
	initModule_nonfree();

	Ptr<FeatureDetector> keypointDetector = FeatureDetector::create("SIFT");

	Ptr<DescriptorExtractor> descriptorExtractor = DescriptorExtractor::create("SIFT");

	Ptr<DescriptorMatcher> matcher = DescriptorMatcher::create("FlannBased");

	BOWImgDescriptorExtractor bowDE(keypointDetector, matcher);

	vector<vector<KeyPoint>> keypoints;

	keypointDetector->detect(imgs, keypoints);

	vector<Mat> descriptors;

	descriptorExtractor->compute(imgs, keypoints, descriptors);

	for (int i = 0; i < NUMBER_IMAGES; i++)
	{
		FileStorage fs(XML_PATH + FILE_LIST[i] + ".xml", FileStorage::WRITE);
		
		fs << "Image" << FILE_LIST[i] + ".jpg";

		fs << "KeyPoints" << keypoints[i];

		fs << "Descriptors" << descriptors[i];

		fs.release();
	}

	//BEGIN: huấn luyện
	TermCriteria tc(TermCriteria::MAX_ITER + TermCriteria::EPS, 10, 0.001);

	int retries = 1;

	int flags = KMEANS_PP_CENTERS;

	BOWKMeansTrainer bowTrainer(DICTIONARY_SIZE, tc, retries, flags);

	for (int i = 0; i < NUMBER_IMAGES; i++)
	{
		bowTrainer.add(descriptors[i]);
	}

	Mat dictionary = bowTrainer.cluster();

	bowDE.setVocabulary(dictionary);
	//END: huấn luyện

	FileStorage fs(VOCABULARY_PATH, FileStorage::WRITE);

	fs << "Vocabulary" << bowDE.getVocabulary();

	fs.release(); 
	
	//BEGIN: Tính toán histogram cho từng ảnh và lưu vào file XMl của ảnh tương ứng
	for (int i = 0; i < NUMBER_IMAGES; i++)
	{
		String abc = XML_PATH + FILE_LIST[i] + ".xml";
		FileStorage fs(XML_PATH + FILE_LIST[i] + ".xml", FileStorage::APPEND);
		
		fs << "Histogram" << CalcHistogram(descriptors[i], bowDE.getVocabulary());

		fs.release();
	}
	//END: Tính toán histogram cho từng ảnh và lưu vào file XMl của ảnh tương ứng
}

/*
Tính histogram cho 1 ảnh trong bộ dữ liệu huấn luyện
*/
Mat CProcessor::CalcHistogram(Mat descriptor, Mat vocabulary)
{
	Mat histogram = Mat::zeros(1, DICTIONARY_SIZE, CV_32F);

	float threshold = 0.9, difference;

	for(int i = 0; i < vocabulary.rows; i++)//duyệt qua các từ trong BoF (tâm các cluster)
	{
		for(int j = 0; j < descriptor.rows; j++)//duyệt qua các descriptor (128D-SIFT)
		{
			difference = compareHist(vocabulary.row(i), descriptor.row(j), CV_COMP_CORREL);//diffrence thuộc [0, 1]:  càng cao thì càng giống nhau

			histogram.at<float>(Point(i, 0)) += 
				difference > threshold
				? difference
				: 0;
		}
	}
	normalize( histogram, histogram, 0, 1, NORM_MINMAX, -1, Mat() );

	return histogram;
}

vector<PAIR> CProcessor::Matching(Mat query)
{
	//BEGIN: Load dữ liệu xml

	vector<string> ImageName;

	vector<vector<KeyPoint>> keypoints;

	vector<Mat> descriptors;

	vector<Mat> histograms;

	for (int i = 0; i < NUMBER_IMAGES; i++)
	{
		FileStorage fs(XML_PATH + FILE_LIST[i] + ".xml", FileStorage::READ);

		vector<KeyPoint> keypoint;

		Mat descriptor, histogram;

		string image;

		fs["Image"] >> image;
		ImageName.push_back(image);

		read(fs["KeyPoints"], keypoint);
		keypoints.push_back(keypoint);

		read(fs["Descriptors"], descriptor);
		descriptors.push_back(descriptor);

		read(fs["Histogram"], histogram);
		histograms.push_back(histogram);

		fs.release();
	}
	//END: Load dữ liệu xml
	
	//BEGIN: Tìm keypoints, descriptors cho ảnh truy vấn
	initModule_nonfree();

	Ptr<FeatureDetector> keypointDetector = FeatureDetector::create("SIFT");

	Ptr<DescriptorExtractor> descriptorExtractor = DescriptorExtractor::create("SIFT");

	Ptr<DescriptorMatcher> matcher = DescriptorMatcher::create("FlannBased");

	BOWImgDescriptorExtractor bowDE(keypointDetector, matcher);

	vector<KeyPoint> queryKeyPoint;

	keypointDetector->detect(query, queryKeyPoint);

	Mat queryDescriptor;

	descriptorExtractor->compute(query, queryKeyPoint, queryDescriptor);
	//END: Tìm keypoints, descriptors cho ảnh truy vấn

	//BEGIN: Tính histogram cho ảnh truy vấn
	Mat Vocabulary;

	FileStorage fs(VOCABULARY_PATH, FileStorage::READ);

	fs["Vocabulary"] >> Vocabulary;

	fs.release();

	Mat queryHist = CalcHistogram(queryDescriptor, Vocabulary);
	//END: Tính histogram cho ảnh truy vấn

	//BEGIN: Matching
	vector<PAIR> Result;

	PAIR pair;

	float difference = 0;

	for (int i = 0; i < NUMBER_IMAGES; i++)
	{
		difference = compareHist(queryHist, histograms[i], CV_COMP_CORREL);

		if (difference > 0.8)
		{
			pair.index = i;
			pair.value = difference;
			Result.push_back(pair);
		}
	}

	for (int i = 0; i < Result.size() - 1; i++)
	{
		for (int j = i + 1; j < Result.size(); j++)
		{
			if (Result[i].value < Result[j].value)
			{
				swap(Result[i], Result[j]);
			}
		}
	}
	//END: Matching

	return Result;
}

void CProcessor::ShowResult(vector<PAIR> result)
{
	for (int i = 0; i < result.size(); i++)
	{
		cout << "Image: " << FILE_LIST[result[i].index] << " - " << "Accuracy: " << result[i].value << endl;
	}
}

vector<string> CProcessor::GetFileList(string FileListPath)
{
	vector<string> list;

	ifstream f;

	f.open(FileListPath, ifstream::in);

	char s[255];

	while (!f.eof())
	{
		f.getline(s, 255);
		list.push_back(s);
	}

	f.close();

	return list;
}